// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2012 Desire Nuentsa <desire.nuentsa_wakam@inria.fr>
// Copyright (C) 2014 Gael Guennebaud <gael.guennebaud@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_SUITESPARSEQRSUPPORT_H
#define EIGEN_SUITESPARSEQRSUPPORT_H

namespace Eigen {

template <typename MatrixType> class SPQR;
template <typename SPQRType> struct SPQRMatrixQReturnType;
template <typename SPQRType> struct SPQRMatrixQTransposeReturnType;
template <typename SPQRType, typename Derived> struct SPQR_QProduct;
namespace internal {
    template <typename SPQRType> struct traits<SPQRMatrixQReturnType<SPQRType>>
    {
        typedef typename SPQRType::MatrixType ReturnType;
    };
    template <typename SPQRType> struct traits<SPQRMatrixQTransposeReturnType<SPQRType>>
    {
        typedef typename SPQRType::MatrixType ReturnType;
    };
    template <typename SPQRType, typename Derived> struct traits<SPQR_QProduct<SPQRType, Derived>>
    {
        typedef typename Derived::PlainObject ReturnType;
    };
}  // End namespace internal

/**
  * \ingroup SPQRSupport_Module
  * \class SPQR
  * \brief Sparse QR factorization based on SuiteSparseQR library
  *
  * This class is used to perform a multithreaded and multifrontal rank-revealing QR decomposition
  * of sparse matrices. The result is then used to solve linear leasts_square systems.
  * Clearly, a QR factorization is returned such that A*P = Q*R where :
  *
  * P is the column permutation. Use colsPermutation() to get it.
  *
  * Q is the orthogonal matrix represented as Householder reflectors.
  * Use matrixQ() to get an expression and matrixQ().transpose() to get the transpose.
  * You can then apply it to a vector.
  *
  * R is the sparse triangular factor. Use matrixQR() to get it as SparseMatrix.
  * NOTE : The Index type of R is always SuiteSparse_long. You can get it with SPQR::Index
  *
  * \tparam _MatrixType The type of the sparse matrix A, must be a column-major SparseMatrix<>
  *
  * \implsparsesolverconcept
  *
  *
  */
template <typename _MatrixType> class SPQR : public SparseSolverBase<SPQR<_MatrixType>>
{
protected:
    typedef SparseSolverBase<SPQR<_MatrixType>> Base;
    using Base::m_isInitialized;

public:
    typedef typename _MatrixType::Scalar Scalar;
    typedef typename _MatrixType::RealScalar RealScalar;
    typedef SuiteSparse_long StorageIndex;
    typedef SparseMatrix<Scalar, ColMajor, StorageIndex> MatrixType;
    typedef Map<PermutationMatrix<Dynamic, Dynamic, StorageIndex>> PermutationType;
    enum
    {
        ColsAtCompileTime = Dynamic,
        MaxColsAtCompileTime = Dynamic
    };

public:
    SPQR()
        : m_analysisIsOk(false), m_factorizationIsOk(false), m_isRUpToDate(false), m_ordering(SPQR_ORDERING_DEFAULT), m_allow_tol(SPQR_DEFAULT_TOL),
          m_tolerance(NumTraits<Scalar>::epsilon()), m_cR(0), m_E(0), m_H(0), m_HPinv(0), m_HTau(0), m_useDefaultThreshold(true)
    {
        cholmod_l_start(&m_cc);
    }

    explicit SPQR(const _MatrixType& matrix)
        : m_analysisIsOk(false), m_factorizationIsOk(false), m_isRUpToDate(false), m_ordering(SPQR_ORDERING_DEFAULT), m_allow_tol(SPQR_DEFAULT_TOL),
          m_tolerance(NumTraits<Scalar>::epsilon()), m_cR(0), m_E(0), m_H(0), m_HPinv(0), m_HTau(0), m_useDefaultThreshold(true)
    {
        cholmod_l_start(&m_cc);
        compute(matrix);
    }

    ~SPQR()
    {
        SPQR_free();
        cholmod_l_finish(&m_cc);
    }
    void SPQR_free()
    {
        cholmod_l_free_sparse(&m_H, &m_cc);
        cholmod_l_free_sparse(&m_cR, &m_cc);
        cholmod_l_free_dense(&m_HTau, &m_cc);
        std::free(m_E);
        std::free(m_HPinv);
    }

    void compute(const _MatrixType& matrix)
    {
        if (m_isInitialized)
            SPQR_free();

        MatrixType mat(matrix);

        /* Compute the default threshold as in MatLab, see:
       * Tim Davis, "Algorithm 915, SuiteSparseQR: Multifrontal Multithreaded Rank-Revealing
       * Sparse QR Factorization, ACM Trans. on Math. Soft. 38(1), 2011, Page 8:3 
       */
        RealScalar pivotThreshold = m_tolerance;
        if (m_useDefaultThreshold)
        {
            RealScalar max2Norm = 0.0;
            for (int j = 0; j < mat.cols(); j++) max2Norm = numext::maxi(max2Norm, mat.col(j).norm());
            if (max2Norm == RealScalar(0))
                max2Norm = RealScalar(1);
            pivotThreshold = 20 * (mat.rows() + mat.cols()) * max2Norm * NumTraits<RealScalar>::epsilon();
        }
        cholmod_sparse A;
        A = viewAsCholmod(mat);
        m_rows = matrix.rows();
        Index col = matrix.cols();
        m_rank = SuiteSparseQR<Scalar>(m_ordering, pivotThreshold, col, &A, &m_cR, &m_E, &m_H, &m_HPinv, &m_HTau, &m_cc);

        if (!m_cR)
        {
            m_info = NumericalIssue;
            m_isInitialized = false;
            return;
        }
        m_info = Success;
        m_isInitialized = true;
        m_isRUpToDate = false;
    }
    /** 
     * Get the number of rows of the input matrix and the Q matrix
     */
    inline Index rows() const { return m_rows; }

    /** 
     * Get the number of columns of the input matrix. 
     */
    inline Index cols() const { return m_cR->ncol; }

    template <typename Rhs, typename Dest> void _solve_impl(const MatrixBase<Rhs>& b, MatrixBase<Dest>& dest) const
    {
        eigen_assert(m_isInitialized && " The QR factorization should be computed first, call compute()");
        eigen_assert(b.cols() == 1 && "This method is for vectors only");

        //Compute Q^T * b
        typename Dest::PlainObject y, y2;
        y = matrixQ().transpose() * b;

        // Solves with the triangular matrix R
        Index rk = this->rank();
        y2 = y;
        y.resize((std::max)(cols(), Index(y.rows())), y.cols());
        y.topRows(rk) = this->matrixR().topLeftCorner(rk, rk).template triangularView<Upper>().solve(y2.topRows(rk));

        // Apply the column permutation
        // colsPermutation() performs a copy of the permutation,
        // so let's apply it manually:
        for (Index i = 0; i < rk; ++i) dest.row(m_E[i]) = y.row(i);
        for (Index i = rk; i < cols(); ++i) dest.row(m_E[i]).setZero();

        //       y.bottomRows(y.rows()-rk).setZero();
        //       dest = colsPermutation() * y.topRows(cols());

        m_info = Success;
    }

    /** \returns the sparse triangular factor R. It is a sparse matrix
     */
    const MatrixType matrixR() const
    {
        eigen_assert(m_isInitialized && " The QR factorization should be computed first, call compute()");
        if (!m_isRUpToDate)
        {
            m_R = viewAsEigen<Scalar, ColMajor, typename MatrixType::StorageIndex>(*m_cR);
            m_isRUpToDate = true;
        }
        return m_R;
    }
    /// Get an expression of the matrix Q
    SPQRMatrixQReturnType<SPQR> matrixQ() const { return SPQRMatrixQReturnType<SPQR>(*this); }
    /// Get the permutation that was applied to columns of A
    PermutationType colsPermutation() const
    {
        eigen_assert(m_isInitialized && "Decomposition is not initialized.");
        return PermutationType(m_E, m_cR->ncol);
    }
    /**
     * Gets the rank of the matrix. 
     * It should be equal to matrixQR().cols if the matrix is full-rank
     */
    Index rank() const
    {
        eigen_assert(m_isInitialized && "Decomposition is not initialized.");
        return m_cc.SPQR_istat[4];
    }
    /// Set the fill-reducing ordering method to be used
    void setSPQROrdering(int ord) { m_ordering = ord; }
    /// Set the tolerance tol to treat columns with 2-norm < =tol as zero
    void setPivotThreshold(const RealScalar& tol)
    {
        m_useDefaultThreshold = false;
        m_tolerance = tol;
    }

    /** \returns a pointer to the SPQR workspace */
    cholmod_common* cholmodCommon() const { return &m_cc; }

    /** \brief Reports whether previous computation was successful.
      *
      * \returns \c Success if computation was successful,
      *          \c NumericalIssue if the sparse QR can not be computed
      */
    ComputationInfo info() const
    {
        eigen_assert(m_isInitialized && "Decomposition is not initialized.");
        return m_info;
    }

protected:
    bool m_analysisIsOk;
    bool m_factorizationIsOk;
    mutable bool m_isRUpToDate;
    mutable ComputationInfo m_info;
    int m_ordering;                 // Ordering method to use, see SPQR's manual
    int m_allow_tol;                // Allow to use some tolerance during numerical factorization.
    RealScalar m_tolerance;         // treat columns with 2-norm below this tolerance as zero
    mutable cholmod_sparse* m_cR;   // The sparse R factor in cholmod format
    mutable MatrixType m_R;         // The sparse matrix R in Eigen format
    mutable StorageIndex* m_E;      // The permutation applied to columns
    mutable cholmod_sparse* m_H;    //The householder vectors
    mutable StorageIndex* m_HPinv;  // The row permutation of H
    mutable cholmod_dense* m_HTau;  // The Householder coefficients
    mutable Index m_rank;           // The rank of the matrix
    mutable cholmod_common m_cc;    // Workspace and parameters
    bool m_useDefaultThreshold;     // Use default threshold
    Index m_rows;
    template <typename, typename> friend struct SPQR_QProduct;
};

template <typename SPQRType, typename Derived> struct SPQR_QProduct : ReturnByValue<SPQR_QProduct<SPQRType, Derived>>
{
    typedef typename SPQRType::Scalar Scalar;
    typedef typename SPQRType::StorageIndex StorageIndex;
    //Define the constructor to get reference to argument types
    SPQR_QProduct(const SPQRType& spqr, const Derived& other, bool transpose) : m_spqr(spqr), m_other(other), m_transpose(transpose) {}

    inline Index rows() const { return m_transpose ? m_spqr.rows() : m_spqr.cols(); }
    inline Index cols() const { return m_other.cols(); }
    // Assign to a vector
    template <typename ResType> void evalTo(ResType& res) const
    {
        cholmod_dense y_cd;
        cholmod_dense* x_cd;
        int method = m_transpose ? SPQR_QTX : SPQR_QX;
        cholmod_common* cc = m_spqr.cholmodCommon();
        y_cd = viewAsCholmod(m_other.const_cast_derived());
        x_cd = SuiteSparseQR_qmult<Scalar>(method, m_spqr.m_H, m_spqr.m_HTau, m_spqr.m_HPinv, &y_cd, cc);
        res = Matrix<Scalar, ResType::RowsAtCompileTime, ResType::ColsAtCompileTime>::Map(reinterpret_cast<Scalar*>(x_cd->x), x_cd->nrow, x_cd->ncol);
        cholmod_l_free_dense(&x_cd, cc);
    }
    const SPQRType& m_spqr;
    const Derived& m_other;
    bool m_transpose;
};
template <typename SPQRType> struct SPQRMatrixQReturnType
{
    SPQRMatrixQReturnType(const SPQRType& spqr) : m_spqr(spqr) {}
    template <typename Derived> SPQR_QProduct<SPQRType, Derived> operator*(const MatrixBase<Derived>& other)
    {
        return SPQR_QProduct<SPQRType, Derived>(m_spqr, other.derived(), false);
    }
    SPQRMatrixQTransposeReturnType<SPQRType> adjoint() const { return SPQRMatrixQTransposeReturnType<SPQRType>(m_spqr); }
    // To use for operations with the transpose of Q
    SPQRMatrixQTransposeReturnType<SPQRType> transpose() const { return SPQRMatrixQTransposeReturnType<SPQRType>(m_spqr); }
    const SPQRType& m_spqr;
};

template <typename SPQRType> struct SPQRMatrixQTransposeReturnType
{
    SPQRMatrixQTransposeReturnType(const SPQRType& spqr) : m_spqr(spqr) {}
    template <typename Derived> SPQR_QProduct<SPQRType, Derived> operator*(const MatrixBase<Derived>& other)
    {
        return SPQR_QProduct<SPQRType, Derived>(m_spqr, other.derived(), true);
    }
    const SPQRType& m_spqr;
};

}  // End namespace Eigen
#endif
